#!/usr/bin/env python
import sys
import os
# 获取当前脚本所在目录（scripts/）
current_dir = os.path.dirname(os.path.abspath(__file__))
# 将auction_python目录（scripts/auction_python/）添加到模块搜索路径
auction_python_dir = os.path.join(current_dir, 'auction_python')
sys.path.append(auction_python_dir)  # 添加模块路径
import rospy
import numpy as np
import time
from geometry_msgs.msg import PoseStamped
from tf.transformations import euler_from_quaternion
from bid_functions import InterceptionPro, bid, bidding

class TaskAssigner:
    def __init__(self):
        # 初始化节点
        rospy.init_node('task_assigner', anonymous=True)
        
        # 参数设置（可通过launch文件配置）
        our_uav_ids_str = rospy.get_param('~our_uav_ids', [1, 4])  # 我方无人机ID字符串列表
        self.our_uav_ids = [int(uid.strip()) for uid in our_uav_ids_str.split(',')]  # 解析为整数列表[1,4]
        enemy_uav_ids_str = rospy.get_param('~enemy_uav_ids', [2, 3])  # 敌方无人机ID字符串列表
        self.enemy_uav_ids = [int(uid.strip()) for uid in enemy_uav_ids_str.split(',')]
        self.update_rate = rospy.get_param('~update_rate', 1.0)  # 分配更新频率(Hz)
        self.detection_range = rospy.get_param('~detection_range', 100.0)  # 探测范围
        self.our_speed = rospy.get_param('~our_speed', 1100.0)  # 我方速度
        self.enemy_speed = rospy.get_param('~enemy_speed', 500.0)  # 敌方速度
        self.enemy_importance = rospy.get_param('~enemy_importance', 1)  # 敌方重要程度
        
        # 存储无人机信息
        self.our_uav_info = {uid: None for uid in self.our_uav_ids}  # {id: [x,y,z,yaw]}
        self.enemy_uav_info = {uid: None for uid in self.enemy_uav_ids}  # {id: [x,y,z,yaw]}
        
        # 订阅所有无人机的位置信息
        self.subscribers = []
        for uid in self.our_uav_ids + self.enemy_uav_ids:
            topic = f'/uav{uid}/mavros/local_position/pose'
            sub = rospy.Subscriber(topic, PoseStamped, self.pose_callback, callback_args=uid)
            self.subscribers.append(sub)
        
        # 初始化分配结果
        self.assignment = {uid: self.enemy_uav_ids[i] for i, uid in enumerate(self.our_uav_ids)}
        self.publish_assignment()
        
        rospy.loginfo("任务分配节点初始化完成")

    def pose_callback(self, msg, uid):
        """回调函数：更新无人机位置和航向角信息"""
        # 提取位置
        x = msg.pose.position.x
        y = msg.pose.position.y
        z = msg.pose.position.z
        
        # 提取航向角（偏航角）
        quat = [msg.pose.orientation.x, msg.pose.orientation.y,
                msg.pose.orientation.z, msg.pose.orientation.w]
        roll, pitch, yaw = euler_from_quaternion(quat)
        
        # 存储信息（弧度转角度，根据原算法需求）
        if uid in self.our_uav_ids:
            self.our_uav_info[uid] = [x, y, z, np.degrees(yaw)]
        else:
            self.enemy_uav_info[uid] = [x, y, z, np.degrees(yaw)]

    def prepare_inputs(self):
        """准备算法输入数据"""
        # 检查是否所有无人机信息都已获取
        if None in self.our_uav_info.values() or None in self.enemy_uav_info.values():
            rospy.logwarn("部分无人机信息未获取，跳过本次分配")
            return None, None
        
        # 构建我方信息矩阵 [x, y, z, 航向角, 探测范围, 速度]
        info_m = []
        for uid in sorted(self.our_uav_ids):  # 按ID排序，确保顺序一致
            x, y, z, yaw = self.our_uav_info[uid]
            info_m.append([x, y, z, yaw, self.detection_range, self.our_speed])
        
        # 构建敌方信息矩阵 [x, y, z, 重要程度, 速度, 航向角]
        info_t = []
        for uid in sorted(self.enemy_uav_ids):  # 按ID排序，确保顺序一致
            x, y, z, yaw = self.enemy_uav_info[uid]
            info_t.append([x, y, z, self.enemy_importance, self.enemy_speed, yaw])
        
        return np.array(info_m), np.array(info_t)

    def update_assignment(self):
        """运行分配算法并更新目标ID参数"""
        info_m, info_t = self.prepare_inputs()
        if info_m is None or info_t is None:
            return
        
        # 调用分配算法
        try:
            result_fenpei = TargetAssign(info_m, info_t)
            rospy.loginfo(f"分配结果: {result_fenpei}")
            
            # 映射分配结果到实际敌方ID（算法返回的是索引+1，需转换为实际ID）
            enemy_sorted = sorted(self.enemy_uav_ids)  # 与info_t顺序一致
            for i, our_uid in enumerate(sorted(self.our_uav_ids)):
                target_idx = result_fenpei[i] - 1  # 转换为0基索引
                self.assignment[our_uid] = enemy_sorted[target_idx]
            
            # 发布分配结果（更新ROS参数）
            self.publish_assignment()
            
        except Exception as e:
            rospy.logerr(f"分配算法执行失败: {str(e)}")

    def publish_assignment(self):
        """通过ROS参数发布目标ID，供我方无人机节点读取"""
        for our_uid, target_uid in self.assignment.items():
            param_name = f'/uav{our_uid}/target_uav_id' 
            rospy.set_param(param_name, target_uid)
            rospy.loginfo(f"已更新uav{our_uid}的目标为uav{target_uid}")

    def run(self):
        """主循环：定期更新分配结果"""
        rate = rospy.Rate(self.update_rate)#以固定频率更新id绑定关系
        while not rospy.is_shutdown():
            self.update_assignment()
            rate.sleep()


def TargetAssign(info_m, info_t):
    # 保持原算法不变，此处复用用户提供的TargetAssign函数
    nm = info_m.shape[0]
    nt = info_t.shape[0]
    num = 1 # 分配次数
    result = np.zeros((num, nm))
    f_gen = np.zeros(num)

    start_time = time.time()
    seed = 2
    np.random.seed(seed)
    x1 = np.array([])
    y1 = np.array([])

    for gen in range(num):
        r = range(nm)  # 适配2架无人机的情况
        info_m1 = info_m[r, :]
        
        P = InterceptionPro(info_m1, info_t)
        F = bid(nm, nt, P, info_t)
        epsilon = 0.1
        H, X, f, x1, y1 = bidding(nt, nm, F, epsilon, x1, y1)
        
        f_gen[gen] = f
        for i in range(nm):
            result[gen, r[i]] = H[i]

    end_time = time.time()
    tt = end_time - start_time

    sorted_indices = np.argsort(f_gen)[::-1]
    result2 = result[sorted_indices[0], :]
    result_fenpei = list(int(float(x)) for x in result2)
    return result_fenpei


if __name__ == '__main__':
    try:
        assigner = TaskAssigner()
        assigner.run()
    except rospy.ROSInterruptException:
        pass
